import warnings
warnings.filterwarnings('ignore')


# 加载飞桨和相关类库
import paddle
import numpy as np

from paddle import nn
from paddle.io import Dataset

from sklearn.model_selection import train_test_split

print(paddle.__version__)

BATCH_SIZE = 32
EPOCH = 10
TEST_SIZE = 0.2
LEARNING_RATE = 0.0001


# 定义网络块，构建神经网络的基本结构
class Block(nn.Layer):
    print("..Block..")
    """
    定义一个基本的网络块，包含卷积、批量归一化、LeakyReLU激活函数和Dropout层
        nn.Sequential 按顺序组合神经网络层
            : nn.Conv1D 创建一个一维卷积层
            : nn.BatchNorm1D 创建一个一维批量归一化层 num_features 输入特征的数量
              （通道数）
            : nn.LeakyReLU 创建一个带有泄露整流线性单元（Leaky ReLU）激活函数的层 
              negative_slope：负斜率系数，默认为0.01
            : nn.Dropout 创建一个用于随机失活的层p为失活概率
        forward 正向传播函数，nn.Sequential
    """
    def __init__(self, in_channels=3, out_channels=3, kernel_size=3, *args):
        super(Block, self).__init__()
        self.nn = nn.Sequential(
            nn.Conv1D(
                in_channels=in_channels,    # 输入通道
                out_channels=out_channels,  # 输出通道
                kernel_size=kernel_size,    # 卷积核大小
                padding=0,                  # 填充大小
                bias_attr=False             # 是否包含偏置参数
                ),
            nn.BatchNorm1D(num_features=out_channels),
            nn.LeakyReLU(.2),
            nn.Dropout(p=.2)
        )

    def forward(self, x):
        return self.nn(x)


class FallNet(nn.Layer):
    print("..FallNet..")
    """
    定义FallNet网络结构
        self.cnn 定义 4 个卷积层
            : cnn0 Block( 3, 64, 7, 0)
            : cnn1 Block( 64, 64, 5, 0)
            : cnn2 Block( 64, 64, 3, 0)
            : cnn3 Block( 64, 64, 1, 0)
        self.avg 自适应平均池化层，根据给定的输出尺寸自动进行调整，不依赖于输入的具体尺寸
        self.rnn 定义循环神经网络层
            : rnn0 = nn.GRU( 145, 64, 1, 0.2)
            : rnn1 Block( 64, 64, 1, 0)
            : rnn2 Block( 64, 4, 3, 0)
        self.cls 分类层 nn.Sequential 定义组合神经网络层
            : nn.Linear 定义全连接层，in_f 1016 out_f 128 对输入数据进行线性变换
            : nn.Dropout 随机失活层，防止过拟合，以0.2的概率将输入部分元素置零，减少神经网络复杂度
            : nn.Linear 全连接层，将 128 种 out_f 转换为动作分类任务的类别数量
            : nn.Softmax 激活函数，将原始输出转换为概率分布
        forward 正向传播函数
            : paddle.concat 沿指定轴 axis -1 连接（拼接）多个张量
            : paddle.flatten 将输入张量展平为一维张量
    """
    def __init__(self, in_channels=3, out_classes=5, hid=64, num=64):
        super(FallNet, self).__init__()

        self.cnn0 = Block(in_channels,hid,7,0)
        self.cnn1 = Block(hid,hid,5,0)
        self.cnn2 = Block(hid,hid,3,0)
        self.cnn3 = Block(hid,hid,1,0)

        self.avg = nn.AdaptiveAvgPool1D(output_size=num)

        self.rnn0 = nn.GRU(input_size=145, hidden_size=num, num_layers=1, dropout=0.2)
        self.rnn1 = Block(hid,hid,1,0)
        self.rnn2 = Block(hid,4,3,0)

        self.cls = nn.Sequential(
            nn.Linear(in_features=1016, out_features=128),
            nn.Dropout(p=.2),
            nn.Linear(in_features=128, out_features=out_classes),
            nn.Softmax(axis=1)
        )

    def forward(self, x):
        x = self.cnn0(x)
        y = self.cnn1(x)
        y1 = self.avg(y)

        y = self.cnn2(y)
        y2 = self.avg(y)

        y = self.cnn3(y)
        y3 = self.avg(y)

        r,t = self.rnn0(x)

        x = paddle.concat([y1,y2,y3,r], axis=-1)

        x = self.rnn1(x)
        x = self.rnn2(x)
        x = paddle.flatten(x, start_axis=1)

        x = self.cls(x)
        return x


class HarDataset(Dataset):
    print("..HarDataset..")
    def __init__(self, mode='train'):
        """
        步骤一：继承paddle.io.Dataset类
        步骤二：实现构造函数，定义数据集大小
            :加载加速度数据集 acc:12341*453 ，加载动作分类标签 label:12341*1
            : mode 分别为训练模式 'train' 和测试模式 'test'
        步骤三：实现__getitem__方法，定义指定index时如何获取数据，并返回单条数据（训练数据，对应的标签）
        步骤四：实现__len__方法，返回数据集规模
        """
        super(HarDataset, self).__init__()

        acc = np.load('./data/db5_acc.npy').reshape(-1, 3, 151)
        label = np.load('./data/db5_lab.npy').astype(np.int64)

        # 随即划分训练集测试集，测试集占比0.3
        self.acc_train, self.acc_test, self.label_train, self.label_test = \
        train_test_split(acc, label, stratify=label, test_size=TEST_SIZE)

        if mode == 'train':
            self.data = self.acc_train, self.label_train
        else:
            self.data = self.acc_test, self.label_test

    def __getitem__(self, index):
        return self.data[0][index], self.data[1][index]

    def __len__(self):
        self.num_scale = self.data[0].shape[0]

        return self.num_scale
    

# 定义并初始化数据读取器
def HarDataset_paddle_loader(mode):
    print("..HarDataset_paddle_loader..")
    return paddle.io.DataLoader(HarDataset(mode), batch_size=BATCH_SIZE, shuffle=True)


# 模型训练
def train(model, dataLoader=None, *args, **kwargs):
    print("..Training..")
    """ model.prepare
            prepare为模型训练做准备，设置优化器及其学习率，并将网络的参数传入优化器，设置损失函数和精度计算方式
            Adam: 使用 Adam 优化器，学习率为 0.0001 传入封装好的全部模型参数 model.parameters 用于后续更新
            CrossEntropyLoss: 使用交叉熵损失函数 CrossEntropyLoss 用于分类任务评估
            Accuracy: 使用分类任务常用的准确率指标 Accuracy 计算模型在训练集上的精度
        model.fit
            启动模型训练，指定训练数据集，设置训练轮次，设置每次数据集计算的批次大小，设置日志格式
            使用 harset_loader 进行数据加载，其中 eval=False 表示训练模式
    """
    if dataLoader == None:
        print("The data loader is not set!")
        return
    
    model.prepare(
        paddle.optimizer.Adam(LEARNING_RATE, parameters=model.parameters()),
        paddle.nn.CrossEntropyLoss(),
        paddle.metric.Accuracy()
    )

    model.fit(
        dataLoader(*args, **kwargs),
        epochs=EPOCH,               # 训练时遍历数据集的次数，即外循环轮次
        batch_size=BATCH_SIZE,      # 内循环中每个批次的训练样本数
        verbose=1,
    )


# 模型保存
def save(model, path):
    print("..Saving..")
    """
    将模型保存为可部署的推理模型（静态图），静态图模型是指一种不包含 Python 运行时的序列化模型，可以在各种平台和设备上高效地部署和运行。
        :存储的模型结构 Program 文件的后缀为 .pdmodel
        :存储的持久参数变量文件的后缀为 .pdiparams
        :变量描述信息存储至文件 .pdiparams.info（额外的信息将在 fine-tune 训练中使用）
    """
    paddle.jit.save(
        # 要保存的模型层
        layer=model.network,
        # 保存的文件路径
        path=path,
        # 输入规范，包括输入形状和数据类型
        input_spec=[paddle.static.InputSpec(shape=[None,3,151],dtype='float32')]
    )


# 模型评估
def eval(model, dataLoader=None, *args, **kwargs):
    print("..Evaluation..")
    if dataLoader == None:
        print("The data loader is not set!")
        return
    
    eval_res = model.evaluate(dataLoader(*args, **kwargs), verbose=1)
    print(eval_res)


def convert(path):
    print("Converting...")
    # 引用Paddlelite预测库
    from paddlelite.lite import Opt

    # 1. 创建opt实例
    opt = Opt()
    # 2. 指定输入模型地址
    # opt.set_model_dir("C:\\Users\\LHZroz\\Desktop\\Roz\\Des\\FallNet\\net005\\")
    opt.set_model_file(path + "test.pdmodel")
    opt.set_param_file(path + "test.pdiparams")
    # 3. 指定转化类型： arm、x86、opencl、npu
    opt.set_valid_places("arm")
    # 4. 指定模型转化类型： naive_buffer、protobuf
    opt.set_model_type("naive_buffer")
    # 3. 输出模型地址
    opt.set_optimize_out("model_opt")
    # 5. 执行模型优化
    opt.run()


if __name__ == '__main__':
    # 0 adl
    # 1 run
    # 2 still
    # 3 fall
    # 4 sos
    # 使用 paddle.Model 将模型封装为一个实例
    model = paddle.Model(FallNet())
    # print(111)
    model.summary((1,3,151))
    print("???")
    # 开始训练
    train(model, HarDataset_paddle_loader, mode='train')

    # 模型评估
    eval(model, HarDataset_paddle_loader, mode='test')

    # 保存模型 C:\Users\LHZroz\Desktop\Roz\Des\FallNet\test
    save(model, "C:/Users/T-T/Desktop/test")

    convert("C:/Users/T-T/Desktop/test/")

    # import paddlelite.lite as lite
    #
    # a = lite.Opt()
    # # 非 combined 形式
    # a.set_model_dir("C:/Users/T-T/Desktop/test/mobilenet_v1/mobilenet_v1")
    #
    # # conmbined 形式，具体模型和参数名称，请根据实际修改
    # # a.set_model_file("D:\\YOU_MODEL_PATH\\mobilenet_v1\\__model__")
    # # a.set_param_file("D:\\YOU_MODEL_PATH\\mobilenet_v1\\__params__")
    #
    # a.set_optimize_out("mobilenet_v1_opt")
    # a.set_valid_places("x86")
    #
    # a.run()